from pymilvus import MilvusClient, FunctionType, Function, RRFRanker, AnnSearchRequest
from pymilvus import connections, db, Collection, FieldSchema, CollectionSchema, DataType

DENSE_INDEX_TYPE = 'FLAT'  # KNN索引方式
M = 100
EFC = 500
SPARSE_INDEX_TYPE = 'SPARSE_INVERTED_INDEX'


# 创建Milvus集合
def create_collection(collection_name):
    print('Creating Collection ...')
    vector_dim = 1024
    block_max_len = 50000  # 定义文本字段的最大长度
    # 定义集合字段
    fields = [
        FieldSchema(name="vecId", dtype=DataType.VARCHAR, is_primary=True, max_length=100),
        FieldSchema(name="blockDenseEmbeddings", dtype=DataType.FLOAT_VECTOR, dim=vector_dim),
        FieldSchema(name="blockSparseEmbeddings", dtype=DataType.SPARSE_FLOAT_VECTOR),
        FieldSchema(name="block", dtype=DataType.VARCHAR, max_length=block_max_len,
                    enable_analyzer=True, analyzer_params={"type": "chinese"}, enable_match=True),
        FieldSchema(name="posInfo", dtype=DataType.VARCHAR, max_length=block_max_len),

    ]
    schema = CollectionSchema(fields=fields, enable_dynamic_field=False)
    schema.add_function(Function(
        name="block_bm25_emb",
        input_field_names=["block"],  # Input block field
        output_field_names=["blockSparseEmbeddings"],  # Internal mapping sparse vector field
        function_type=FunctionType.BM25,  # Model for processing mapping relationship
    ))
    collection = Collection(name=collection_name, schema=schema)

    # 为密集向量字段创建索引
    collection.create_index("blockDenseEmbeddings", {
        'index_type': DENSE_INDEX_TYPE,
        'metric_type': 'COSINE',
    })
    # 为稀疏向量字段创建索引
    collection.create_index("blockSparseEmbeddings", {
        'index_type': SPARSE_INDEX_TYPE,
        'metric_type': 'BM25',
    })
    print(f'{collection_name}创建成功！')
    return collection


# Milvus操作类
class MilvusOP:
    def __init__(self, db_name, collection_name, host, port):
        super().__init__()
        uri = f'http://{host}:{port}'
        alias = "default"
        # 连接Milvus
        if not connections.has_connection(alias=alias):
            connections.connect(alias=alias, host=host, port=port)
            if db_name not in db.list_database():
                db.create_database(db_name)  # 如果数据库不存在，创建数据库
            db.using_database(db_name)
        # 初始化Milvus客户端
        self.client = MilvusClient(uri=uri, db_name=db_name)
        # 如果集合不存在，创建集合
        if collection_name not in self.client.list_collections():
            self.collection = create_collection(collection_name)
        self.collection = Collection(name=collection_name)
        self.collection.load()  # 加载集合数据
        self.output_fields = ['block', 'posInfo']
        self.ranker = RRFRanker()
        self.collection_name = collection_name

    def create_collection(self):
        if self.collection_name not in self.client.list_collections():
            self.collection = create_collection(self.collection_name)
            self.collection = Collection(name=self.collection_name)
            self.collection.load()  # 加载集合数据

    # 稀疏向量搜索
    def block_sparse_search(self, queries, limit=10, expr=None):
        search_params = {
            "metric_type": 'BM25',
            "params": {"drop_ratio_build": 0.0},
        }
        # 使用BM25编码查询并执行稀疏向量搜索
        hit_res = self.collection.search(
            data=queries,
            anns_field="blockSparseEmbeddings",
            limit=limit,
            output_fields=self.output_fields,
            param=search_params,
            expr=expr,
        )
        return hit_res

    def delete(self, vec_ids):
        if vec_ids:
            return self.client.delete(
                collection_name=self.collection_name,
                ids=vec_ids,
            )

    def get_all_vecIDs(self):
        iterator = self.collection.query_iterator(output_fields=["vecId"])
        vec_ids = set()
        while True:
            res = iterator.next()
            if len(res) == 0:
                iterator.close()
                break
            else:
                for vec in res:
                    vec_ids.add(vec['vecId'])
        return vec_ids

    # 混合搜索
    def hybrid_search(self,
                      queries,
                      limit=10,
                      query_embs=None,  # 语义向量
                      dense_block=True,  # 正文，语义检索
                      sparse_block=True,  # 正文，稀疏检索
                      ):
        if not dense_block and not sparse_block:
            sparse_block = True
        expr = None

        reqs = []
        if dense_block and query_embs is not None:
            reqs.append(AnnSearchRequest(data=query_embs, anns_field="blockDenseEmbeddings", limit=limit, expr=expr,
                                         param={"metric_type": "COSINE"}))
        if sparse_block:
            reqs.append(AnnSearchRequest(data=queries, anns_field="blockSparseEmbeddings", limit=limit, expr=expr,
                                         param={"metric_type": "BM25", "params": {"drop_ratio_build": 0.0}}))
        results = self.collection.hybrid_search(reqs=reqs, rerank=self.ranker, limit=limit,
                                                output_fields=self.output_fields)
        return results
